from QKDNetwork import QKDNetwork
from docplex.mp.model import Model
import random
from copy import deepcopy


# 发射密钥
def emit_photon_keys(routeList: list, net: QKDNetwork, p_normalized):
    tmp_g = net.G
    # 按照 p_normalized 依概率发射向链路发射密钥

    emit_direction = []
    all_keys = [i for i in p_normalized]
    for i in all_keys:
        value_tip = p_normalized[i]
        values = list(value_tip.keys())
        probabilities = list(value_tip.values())
        selected_key = random.choices(values, weights=probabilities)[0]
        emit_direction.append([i, selected_key])

    # 发射密钥到密钥池中
    for i in emit_direction:
        tmp_g.edges[i[0], i[1]]["remain_key_num"] += tmp_g.nodes[i[0]][
            "transmitter_rate"
        ]


def provision_with_greedy(routeList: list, net: QKDNetwork, p_normalized):
    tmp_g = net.G
    emit_photon_keys(routeList, net, p_normalized)

    # 按照 需求量Demand - 已满足的密钥量satisfy_num 从小到大排序，然后按照路径长度从短到长排序
    routeList.sort(key=lambda x: (x["sd"][2] - x["sd"][3], len(x["path"])))

    # 遍历routeList，开始抽密钥
    for i in routeList:
        sd, path = i["sd"], i["path"]
        need_key_num = sd[2] - sd[3]
        if need_key_num <= 0:
            continue
        bottle_neck = 1000000
        for j in range(len(path) - 1):
            u, v = path[j], path[j + 1]
            bottle_neck = min(tmp_g.edges[u, v]["remain_key_num"], bottle_neck)
        satisfy_key_num = min(bottle_neck, need_key_num)
        i["sd"] = (sd[0], sd[1], sd[2], sd[3] + satisfy_key_num)
        for j in range(len(path) - 1):
            u, v = path[j], path[j + 1]
            tmp_g.edges[u, v]["remain_key_num"] -= satisfy_key_num
        for sd_idx, sd_item in enumerate(net.sd_list):
            if sd_item[0] == i["sd"][0] and sd_item[1] == i["sd"][1]:
                net.sd_list[sd_idx] = i["sd"]
